import json
from enum import Enum

import sseclient
import streamlit as st
import requests

BACKEND_URL = "http://127.0.0.1:8001/process"


class MSG_TYPE(Enum):
    THINK = 1
    DONE = 2
    ANSWER = 3
    THINK_CHUNK = 4
    QRCODE = 5
    QXQRCODE = 6


# ========= 必须第一个Streamlit命令 =========
st.set_page_config(page_title="聊天机器人（可自动发小红书）", page_icon="💬", layout="centered")
# =========================================

# 样式
st.markdown("""
<style>
.think-box {
    background-color: #f5f7ff;
    padding: 10px;
    border-radius: 8px;
    white-space: pre-wrap;
    font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas,
                 "Liberation Mono", "Courier New", monospace;
    line-height: 1.4;
    border: 1px solid rgba(0,0,0,0.06);
}
</style>
""", unsafe_allow_html=True)


def run_sse_stream(user_input: str, session_state, response_msg):
    """
    连接 FastAPI SSE 接口，把思考过程、最终回答及图片都显示在同一 assistant 气泡中。
    同时将思考过程与最终结果都写入 session_state.messages 供历史回放。
    """
    payload = {"input": user_input}

    # 当前assistant气泡结构
    with response_msg:
        exp = st.expander("🤔 思考过程（点击展开）", expanded=True)
        with exp:
            think_placeholder = st.empty()  # 思考文本实时刷新
        answer_placeholder = st.empty()  # 最终回答
        image_placeholder = st.empty()  # 图片

    think_text = ""  # 累积思考文本
    answer_text = ""
    final_answer = ""
    picture_path = None

    # 对于 sseclient 0.0.27 版本，需要正确设置 headers 并直接传递响应对象
    headers = {'Content-Type': 'application/json'}
    r = requests.post(BACKEND_URL, stream=True, json=payload, headers=headers)
    client = sseclient.SSEClient(r.raw, headers=r.headers)
    last_msg_type = ""
    for msg in client.events():
        if not msg.data.strip():
            continue

        try:
            data = json.loads(msg.data)
        except json.JSONDecodeError:
            think_text += f"\n⚠️ 无法解析: {msg.data}"
            think_placeholder.markdown(f"<div class='think-box'>{think_text}</div>", unsafe_allow_html=True)
            continue

        msg_type = data.get("type")

        if msg_type == MSG_TYPE.THINK.name or msg_type == MSG_TYPE.THINK_CHUNK.name:
            # 思考事件：实时追加
            content = data.get("content", "")
            if msg_type == MSG_TYPE.THINK.name:
                think_text += ("\n" if think_text else "") + content
            elif msg_type == MSG_TYPE.THINK_CHUNK.name:
                think_text += "\n" if last_msg_type != MSG_TYPE.THINK_CHUNK.name else ""
                think_text += content
                last_msg_type = msg_type
            think_placeholder.markdown(f"<div class='think-box'>{think_text}</div>", unsafe_allow_html=True)
        if msg_type == MSG_TYPE.ANSWER.name:
            # 清空图片
            image_placeholder.empty()
            # 回答字段
            content = data.get("content", "")
            answer_text += content
            answer_placeholder.markdown(answer_text, unsafe_allow_html=True)
        elif msg_type == MSG_TYPE.DONE.name:
            # 清空图片
            image_placeholder.empty()
            # 完成事件
            final_answer = data.get("content", "")
            # 展示最终回答
            answer_placeholder.markdown(final_answer, unsafe_allow_html=True)
            break  # 结束循环
        elif msg_type == MSG_TYPE.QRCODE.name:
            qrcode_path = data.get("content", "")
            if qrcode_path:
                image_placeholder.image(qrcode_path, caption="请扫码登陆", width=300)
        elif msg_type == MSG_TYPE.QXQRCODE.name:
            # 收到取消二维码事件时清空图片
            print("收到二维码事件清空")
            image_placeholder.empty()
            answer_placeholder.text("已经成功登陆小红书")

    # 保存到会话历史（包含思考文本）
    session_state.messages.append({
        "role": "assistant",
        "content": final_answer,
        "think": think_text,
        "picture": picture_path
    })


# ===== 页面主体 =====
st.title("💬 旅行方向自动小红书")
st.write("和智能机器人进行对话")

# 初始化历史
if "messages" not in st.session_state:
    st.session_state.messages = []

# 展示历史
for msg in st.session_state.messages:
    if msg["role"] == "user":
        st.chat_message("user").write(msg["content"])
    else:
        with st.chat_message("assistant"):
            # 历史思考过程（可折叠）
            if msg.get("think"):
                with st.expander("🤔 思考过程（点击展开）", expanded=False):
                    st.markdown(f"<div class='think-box'>{msg['think']}</div>", unsafe_allow_html=True)
            # 最终回答
            st.markdown(msg["content"], unsafe_allow_html=True)

# 输入框
if prompt := st.chat_input("请输入您的问题..."):
    st.chat_message("user").write(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    response_msg = st.chat_message("assistant")
    run_sse_stream(prompt, st.session_state, response_msg)
